import argparse
import multiprocessing as mp
import os
from functools import partial
from time import time as timer

from pytube import YouTube
from tqdm import tqdm
from use_sqlite import SqlHandler, Video
from multiprocessing import Pool, cpu_count


class YTDownloader():
    '''
    YTDownloader
    '''

    def __init__(self, proxy_url, output_dir='./videos', numworkers=8):
        self.num_workers = numworkers
        os.makedirs(output_dir, exist_ok=True)
        self.output_dir = output_dir
        # self.proxy_url = proxy_url
        self.proxies = {'http': 'http://' + proxy_url, 'https': 'http://' + proxy_url}
        self.sql_handler = SqlHandler()
        self.can_download_flag: bool = True

    @staticmethod
    def download_single_video_handler(video_url, output_dir="./videos", proxy: str = ''):
        video_postfix = video_url.split('=')[-1]
        video_path = '%s/%s.mp4' % (output_dir, video_postfix)
        if not os.path.isfile(video_path):
            try:
                # Download the highest quality mp4 stream.
                if proxy == "":
                    yt = YouTube(video_url)
                else:
                    yt = YouTube(video_url, proxies={'http': proxy, 'https': proxy})
                stream = yt.streams.filter(subtype='mp4', only_video=True, adaptive=True).first()
                if stream is None:
                    stream = yt.streams.filter(subtype='mp4').first()
                stream.download(output_path=output_dir, filename=video_postfix + '.mp4')
            except Exception as e:
                print('Failed to download %s' % video_url)
                print(e)
        else:
            print('File exists: %s' % video_url)

    def download_video_by_url(self, video_url):
        '''
        暂时不用于多线程
        :param video_url:
        :return:
        '''
        video_postfix = video_url.split('=')[-1]
        video_path = '%s/%s.mp4' % (self.output_dir, video_postfix)
        if not os.path.isfile(video_path):
            try:
                # Download the highest quality mp4 stream.
                yt = YouTube(video_url, proxies=self.proxies)
                stream = yt.streams.filter(subtype='mp4', only_video=True, adaptive=True).first()
                if stream is None:
                    stream = yt.streams.filter(subtype='mp4').first()
                stream.download(output_path=self.output_dir, filename=video_postfix + '.mp4')
            except Exception as e:
                print('Failed to download %s' % video_url)
                print(e)
        else:
            print('File exists: %s' % video_url)

    def add_urls_list_to_db(self, urls_list):
        '''
        添加一个列表到数据库
        :param urls_list:
        :return:
        '''
        for url in urls_list:
            self.add_url_to_db(url)

    def add_mutil_urls_to_db_and_download(self, mutil_url):
        '''
        添加一串类似是'url1,url2,url3'这种形式的字符串到数据库
        :param mutil_url:
        :return:
        '''
        url_list = mutil_url.split(',')
        for url in url_list:
            self.add_url_to_db(url)
        self.download_video_list()

    def check_download(self, video_uuid, video_url):
        '''
        因为多线程下载是不准确的，我们不能知道是否是真实下载好了，需要有一个函数从文件系统中判断是否已经下载完毕。
        暂定True
        :return:
        '''
        return True

    def add_url_to_db(self, url):
        self.sql_handler.insert_video(url)

    def download_video_list(self, number=64):
        if self.can_download_flag:
            self.can_download_flag = False
            video_uuids_urls = self.sql_handler.get_video_uuid_and_url_list(number=number)
            # 获取若干个url

            video_urls = []

            for index, (video_uuid, url) in enumerate(video_uuids_urls):
                video_urls.append(url)
                self.sql_handler.set_video_download(video_uuid=video_uuid)

            # Download videos.
            downloader = partial(YTDownloader.download_single_video_handler)

            start = timer()
            pool_size = self.num_workers
            print('Using pool size of %d' % (pool_size))
            with Pool(processes=pool_size) as p:
                _ = list(tqdm(p.imap_unordered(downloader, video_urls), total=len(video_urls)))
            print('Elapsed time: %.2f' % (timer() - start))
            for video_uuid, url in video_uuids_urls:
                if self.check_download(video_uuid, video_url=url):
                    self.sql_handler.set_video_downloaded_file(video_uuid=video_uuid)
                    # 将已经下载的视频检查下载文件，设置为已下载
                else:
                    self.sql_handler.set_video_download_error(video_uuid=video_uuid)
                    # 设置为下载错误
            self.can_download_flag = True
            return '下载完了'
        else:
            print('已经在下载进程中')
            return '正在下载，请点击统计按钮查看信息'

    def get_current_progress(self):
        now_total_num = self.sql_handler.get_total_videos_count()
        return now_total_num - self.sql_handler.get_undownloaded_videos_count(), now_total_num
